"""LLM Plays Pokemon AI决策循环模块"""

import time
import logging
import threading
from typing import Dict, Any, List, Optional

from llm_pokemon.knowledge.base import KnowledgeBase
from llm_pokemon.prompt.manager import PromptManager
from llm_pokemon.llm.client import LLMClient

logger = logging.getLogger(__name__)


class AILoop:
    """AI决策循环类，负责组成提示词、调用模型、解析响应、执行工具调用和保存状态"""
    
    def __init__(self, config: Dict[str, Any], shared_state, emulator_tool):
        """初始化AI决策循环
        
        Args:
            config: 系统配置
            shared_state: 线程间共享状态
            emulator_tool: 模拟器工具实例（从主线程传入，不再创建新实例）
        """
        self.config = config
        self.shared_state = shared_state
        self.emulator = emulator_tool
        
        # 初始化知识库
        self.knowledge_base = KnowledgeBase(
            storage_path=config["knowledge_base"]["storage_path"],
            backup_dir=config["knowledge_base"]["backup_dir"],
            max_backups=config["knowledge_base"]["max_backups"]
        )
        
        # 初始化提示管理器
        self.prompt_manager = PromptManager(
            max_turns=config["conversation"]["max_turns"],
            history_path=config["conversation"]["history_path"]
        )
        
        # 导入工具定义
        from llm_pokemon.tools.definitions import get_emulator_only_tools_definition
        
        # 使用工具定义模块中的定义
        self.tools_definition = get_emulator_only_tools_definition()
        
        # 初始化LLM客户端
        self.llm_client = LLMClient(
            provider=config["llm"]["provider"],
            model=config["llm"]["model"],
            api_key_env=config["llm"]["api_key_env"],
            max_tokens=config["llm"]["max_tokens"],
            base_url=config["llm"].get("base_url", ""),
            tools=self.tools_definition
        )
        
        # 工具映射
        self.tools = {}
        
        # 导入并初始化use_emulator工具
        from llm_pokemon.tools.use_emulator import UseEmulatorTool
        use_emulator_tool = UseEmulatorTool(self.shared_state, self.emulator)
        
        # 添加到工具映射
        self.tools["use_emulator"] = use_emulator_tool.execute
    
    def start(self):
        """启动AI决策循环"""
        try:
            logger.info("AI决策循环线程启动...")
            
            # 等待模拟器就绪
            logger.info("等待模拟器就绪...")
            while self.shared_state.running and not self.shared_state.emulator_ready:
                time.sleep(0.1)
            
            if not self.shared_state.running:
                logger.info("系统已停止，AI决策循环退出")
                return
            
            logger.info("模拟器就绪，开始AI决策循环")
            
            # 主循环
            while self.shared_state.running:
                try:
                    # 1. 组成提示词
                    prompt = self.prompt_manager.compose_prompt(self.knowledge_base)
                    
                    # 2. 调用模型
                    try:
                        response = self.llm_client.generate(prompt)
                        logger.info(f"LLM响应: {response}")
                    except Exception as e:
                        logger.error(f"调用LLM模型失败: {e}")
                        response = "继续游戏"
                    
                    # 3. 解析响应中的工具调用
                    tool_calls = self._parse_tool_calls(response)
                    
                    # 4. 执行工具调用
                    tool_results = self._execute_tools(tool_calls)
                    
                    # 5. 更新会话历史
                    self.prompt_manager.add_assistant_message(response if isinstance(response, str) else response.content)
                    if tool_results:
                        self.prompt_manager.add_system_message("\n\n".join(tool_results))
                    
                    # 6. 检查是否需要总结（异步方式）
                    if self.prompt_manager.needs_summarization():
                        self._perform_summarization_async()
                    
                    # 7. 保存状态（异步方式）
                    self._save_state_async()
                    
                    # 等待截图更新
                    time.sleep(0.05)
                    
                    # 8. 添加系统消息记录执行结果
                    self._record_execution_results(tool_calls)
                    
                except KeyboardInterrupt:
                    logger.info("用户中断，停止AI决策循环...")
                    with self.shared_state.lock:
                        self.shared_state.running = False
                except Exception as e:
                    logger.exception(f"AI决策循环迭代发生错误: {e}")
                    self.prompt_manager.add_system_message(f"错误: {str(e)}")
                    time.sleep(0.5)  # 出错后稍微暂停
        
        except KeyboardInterrupt:
            logger.info("用户中断，正在关闭AI决策循环线程...")
            with self.shared_state.lock:
                self.shared_state.running = False
        except Exception as e:
            logger.exception(f"AI决策循环线程运行时发生错误: {e}")
        finally:
            logger.info("AI决策循环线程已关闭")
    
    def _parse_tool_calls(self, response) -> List[Dict[str, Any]]:
        """解析模型响应中的工具调用
        
        Args:
            response: 模型响应
            
        Returns:
            工具调用列表
        """
        tool_calls = []
        
        # 如果响应是字符串，创建一个默认的工具调用
        if isinstance(response, str):
            # 默认按A键继续游戏
            tool_calls.append({
                "name": "use_emulator",
                "arguments": {"buttons": ["A"]}
            })
        # 如果响应是OpenAI的消息对象，并且包含工具调用
        elif hasattr(response, 'tool_calls') and response.tool_calls:
            import json
            for tool_call in response.tool_calls:
                try:
                    # 解析函数调用参数
                    function_name = tool_call.function.name
                    function_args = json.loads(tool_call.function.arguments)
                    
                    tool_calls.append({
                        "name": function_name,
                        "arguments": function_args
                    })
                    
                    logger.info(f"解析到工具调用: {function_name}, 参数: {function_args}")
                except Exception as e:
                    logger.error(f"解析工具调用时发生错误: {e}")
        
        return tool_calls
    
    def _execute_tools(self, tool_calls: List[Dict[str, Any]]) -> List[str]:
        """执行工具调用
        
        Args:
            tool_calls: 工具调用列表
            
        Returns:
            工具执行结果列表
        """
        tool_results = []
        for call in tool_calls:
            tool_name = call.get("name")
            tool_args = call.get("arguments", {})
            
            if tool_name in self.tools:
                try:
                    result = self.tools[tool_name](**tool_args)
                    tool_results.append(f"Tool: {tool_name}\nResult: {result}")
                    logger.info(f"执行工具 {tool_name} 成功: {result}")
                except Exception as e:
                    error_msg = f"执行工具 {tool_name} 时发生错误: {e}"
                    logger.error(error_msg)
                    tool_results.append(f"Tool: {tool_name}\nError: {error_msg}")
            else:
                logger.warning(f"未知工具: {tool_name}")
                tool_results.append(f"Tool: {tool_name}\nError: 未知工具")
        
        return tool_results
    
    def _perform_summarization_async(self):
        """异步执行会话总结"""
        logger.info("执行会话总结...")
        
        def perform_summarization():
            try:
                # 构建总结提示
                summarization_prompt = self.prompt_manager.compose_summarization_prompt(
                    self.prompt_manager.get_history(), self.knowledge_base
                )
                # 调用模型生成总结
                summary = self.llm_client.generate(summarization_prompt)
                # 更新会话历史
                self.prompt_manager.clear_history()
                self.prompt_manager.add_assistant_message(summary if isinstance(summary, str) else summary.content)
                logger.info("会话总结完成")
            except Exception as e:
                logger.error(f"执行会话总结时发生错误: {e}")
        
        # 创建异步总结线程
        summarization_thread = threading.Thread(target=perform_summarization)
        summarization_thread.daemon = True  # 设置为守护线程
        summarization_thread.start()
    
    def _save_state_async(self):
        """异步保存状态"""
        def save_state():
            try:
                self.knowledge_base.save()
                self.prompt_manager.save_history()
            except Exception as e:
                logger.error(f"保存状态时发生错误: {e}")
        
        # 创建异步保存线程
        save_thread = threading.Thread(target=save_state)
        save_thread.daemon = True  # 设置为守护线程
        save_thread.start()
    
    def _record_execution_results(self, tool_calls: List[Dict[str, Any]]):
        """记录执行结果
        
        Args:
            tool_calls: 工具调用列表
        """
        with self.shared_state.lock:
            if self.shared_state.screenshot_path:
                buttons_executed = []
                for call in tool_calls:
                    if call.get("name") == "use_emulator" and "buttons" in call.get("arguments", {}):
                        buttons_executed.extend(call["arguments"]["buttons"])
                
                if buttons_executed:
                    self.prompt_manager.add_system_message(f"执行按钮 {buttons_executed} 成功，截图: {self.shared_state.screenshot_path}")